import os
from pprint import pprint

import numpy as np

from my_io import readcsv
import json


def get_idxes(data_type):
    file_path = os.path.join(datasets_root, dataset_path, 'samples/split_sets', data_type+'.txt')
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip().split(' ') for line in lines]
    lines = [line[0] for line in lines]
    lines = sorted(lines)
    indexes = []
    for line in lines:
        assert line in sweeps
        idx = sweeps.index(line)
        indexes.append(idx)
    return indexes


if __name__ =='__main__':
    datasets_root = '/media/adt/T7/ZWH/docker/files/data/motion/datasets'
    assert os.path.exists(datasets_root), f'{datasets_root} does not exist!'
    dataset_paths = []
    for dataset_path in sorted(os.listdir(datasets_root)):
        if dataset_path.startswith('__'):
            dataset_paths.append(dataset_path)

    for dataset_path in dataset_paths:
        calib_json = os.path.join(datasets_root, dataset_path, 'calib/calib.json')
        with open(calib_json, 'r') as f:
            json_data = json.load(f)


        def replace_base_keys(obj):
            for val in obj.values():
                if 'transforms' in val:
                    if 'base' in val['transforms']:
                        val['transforms']['base_link'] = val['transforms']['base']
                        del val['transforms']['base']


        replace_base_keys(json_data)

        keys = json_data.keys()
        cam_keys = []
        for key in keys:
            if key.startswith('camera'):
                cam_keys.append( key)

        for cam_key in cam_keys:
            json_data[cam_key]['transforms'] = dict()
            json_data[cam_key]['ori'] = dict()
            json_data[cam_key]["normal"]= dict()

            json_data[cam_key]['transforms']['base_link'] = np.linalg.inv(np.array(json_data['base']['transforms'][cam_key])).tolist()
            del json_data['base']
            json_data[cam_key]['ori']['imgh'] = json_data[cam_key]['imgh']
            json_data[cam_key]['ori']['imgw'] = json_data[cam_key]['imgw']
            json_data[cam_key]['ori']['K'] = json_data[cam_key]['K_ori']
            json_data[cam_key]['ori']['D'] = json_data[cam_key]['D']
            json_data[cam_key]["normal"]["imgh"] = json_data[cam_key]['imgh_new']
            json_data[cam_key]["normal"]["imgw"] = json_data[cam_key]['imgw_new']
            json_data[cam_key]["normal"]["K"] = json_data[cam_key]['K']
            del json_data[cam_key]['imgh']
            del json_data[cam_key]['imgw']
            del json_data[cam_key]['K_ori']
            del json_data[cam_key]['D']
            del json_data[cam_key]['imgh_new']
            del json_data[cam_key]['imgw_new']
            del json_data[cam_key]['K']


        with open(calib_json, 'w') as f:
            json.dump(json_data, f, ensure_ascii=False, indent=2)

        pass




